import pyaudio
import wave
from pydub import AudioSegment
import os
import time
import struct  # 用于解析音频数据

# ================= 参数配置 =================
FORMAT = pyaudio.paInt16
CHANNELS = 2
RATE = 44100
CHUNK = 1024
THRESHOLD = 1500        # 用于判断静音的音量阈值
SILENCE_TIMEOUT = 10.0   # 静默超过 10 秒则停止
MAX_RECORDING_TIME = 600  # 最大录制时间（秒）

WAVE_OUTPUT_FILENAME = r"output.wav"
MP3_OUTPUT_FILENAME = r"开始懂了.mp3"
CABLE_DEVICE_NAME = "CABLE Output"  # 或 "VB-Audio Virtual Cable"
# ===========================================

p = pyaudio.PyAudio()

def find_vb_cable_input():
    for i in range(p.get_device_count()):
        dev_info = p.get_device_info_by_index(i)
        if CABLE_DEVICE_NAME in dev_info['name']:
            return i, dev_info['maxInputChannels']
    raise Exception(f"未找到 {CABLE_DEVICE_NAME} 输入设备，请检查音频设置")

device_index, max_channels = find_vb_cable_input()
CHANNELS = min(CHANNELS, max_channels)
print(f"使用设备: {p.get_device_info_by_index(device_index)['name']} (Index: {device_index})")

stream = p.open(format=FORMAT,
                channels=CHANNELS,
                rate=RATE,
                input=True,
                input_device_index=device_index,
                frames_per_buffer=CHUNK)

print("系统将在 3 秒后开始录音...")
time.sleep(3)
print("开始录音！")

frames = []
silence_start = None
last_sound_time = time.time()  # 开始时认为有声音，避免立刻结束

try:
    while True:
        data = stream.read(CHUNK, exception_on_overflow=False)

        # 解析音频数据并计算 RMS（用于判断是否静音）
        samples = struct.unpack_from(f"<{len(data)//2}h", data)
        sum_squares = sum(sample * sample for sample in samples)
        rms = (sum_squares / len(samples)) ** 0.5

        # 判断是否有声音
        if rms > THRESHOLD:
            last_sound_time = time.time()
            silence_start = None
        else:
            if silence_start is None:
                silence_start = time.time()

        # 添加当前帧到录制数据中（无论是否有声音）
        frames.append(data)

        # 检查是否静音过久
        if silence_start and (time.time() - silence_start) > SILENCE_TIMEOUT:
            print(f"静默超过 {SILENCE_TIMEOUT} 秒，停止录制。")
            break

        # 防止录制过长
        if time.time() - last_sound_time > MAX_RECORDING_TIME:
            print("达到最大录制时间，停止录制。")
            break

except KeyboardInterrupt:
    print("\n用户手动中断录制。")

finally:
    stream.stop_stream()
    stream.close()
    p.terminate()

if len(frames) == 0:
    print("未录制到任何音频数据。")
else:
    print("正在保存录音文件...")

    # 保存为 WAV
    wf = wave.open(WAVE_OUTPUT_FILENAME, 'wb')
    wf.setnchannels(CHANNELS)
    wf.setsampwidth(p.get_sample_size(FORMAT))
    wf.setframerate(RATE)
    wf.writeframes(b''.join(frames))
    wf.close()

    # 转为 MP3
    sound = AudioSegment.from_wav(WAVE_OUTPUT_FILENAME)
    sound.export(MP3_OUTPUT_FILENAME, format="mp3")

    # 删除临时 WAV 文件
    os.remove(WAVE_OUTPUT_FILENAME)

    print(f"录音完成，文件已保存为: {MP3_OUTPUT_FILENAME}")